package com.zhentao.springai.controller;

import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.model.ChatResponse;
import java.util.List;

@RestController
@CrossOrigin(origins = "*")
public class DeepSeekController {
    
    @Autowired
    private ChatModel chatModel;
    
    @GetMapping("/ai/generate")
    public String deepSeek(@RequestParam(value = "message", defaultValue = "hello") String message) {
        List<Message> messages = List.of(
                new SystemMessage("你是中医助手。请用专业且通俗的中文提供中医相关帮助，包括望闻问切、辨证论治、食疗与穴位按摩建议。请勿提供西药处方，并在给出建议时附上温和的健康提示。"),
                new UserMessage(message)
        );

        Prompt prompt = new Prompt(messages);
        ChatResponse response = chatModel.call(prompt);
        String content = response.getResult().getOutput().getContent();
        System.out.println(content);
        return content;
    }
    
    @GetMapping(value = "/ai/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> streamGenerate(@RequestParam(value = "message", defaultValue = "hello") String message) {
        List<Message> messages = List.of(
                new SystemMessage("你是中医助手。请用专业且通俗的中文提供中医相关帮助，包括望闻问切、辨证论治、食疗与穴位按摩建议。请勿提供西药处方，并在给出建议时附上温和的健康提示。"),
                new UserMessage(message)
        );
        Prompt prompt = new Prompt(messages);
        
        return chatModel.stream(prompt)
                .map(chatResponse -> {
                    String content = chatResponse.getResult().getOutput().getContent();
                    System.out.println("Streaming: " + content);
                    return "data: " + content + "\n\n";
                })
                .onErrorResume(throwable -> {
                    System.err.println("Stream error: " + throwable.getMessage());
                    return Flux.just("data: [ERROR] " + throwable.getMessage() + "\n\n");
                });
    }
}
